#!/bin/bash
#
# Copyright (c) 2020 NVI, Inc.
#
# This file is part of FSL10 Linux distribution.
# (see http://github.com/nvi-inc/fsl10).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#

# For systems with GPT LVM RAID arrays, copies the partition table onto
# the second disk (if missing), rebuilds the boot sector and adds the RAID
# partition back into the array as a spare to be rebuilt.

usage() { echo "Usage: $0 [-A][-h]"; echo -e "  where\t-A\tAllow \"Secondary\" disk to be loaded at start"; echo -e "  \t-h\tdisplay this help and exit"; }

ACCEPT_LOAD=0
while getopts "Ah" arg; do
	case $arg in
		A)
			ACCEPT_LOAD=1
			;;
		h | *)
			usage
			exit 0
			;;
	esac
done

disk="/dev/sdb"
devname=$(basename $disk)
part1="$disk"1
part2="$disk"2
lockf=/tmp/.accept_load

if [ ! -b /dev/md0 ]; then
	echo "ERROR: RAID /dev/md0 does not exist!!"
	exit 1
fi
if mdadm --detail /dev/md0 2>&1 | grep -qE ": (clean|active) $" ; then
	echo "ERROR: RAID /dev/md0 does not need recovery!!"
	exit 1
fi
if ! mdadm --detail /dev/md0 2>&1 | grep -q ", degraded $" ; then
	echo "ERROR: RAID /dev/md0 is not in a recoverable state"
	exit 1
fi

# Perform sanity checks on primary disk first
if ! mdadm --detail /dev/md0 | grep -q /dev/sda; then
	echo "ERROR: \"Primary\" disk /dev/sda is not part of RAID /dev/md0!!"
	exit 1
fi
if [ "$(cat /sys/block/sda/removable)" != "0" ]; then
	echo "WARNING: \"Primary\" disk /dev/sda is removable:"
	echo -n " /dev/sda: "$(lsblk --nodeps -no tran /dev/sda | tr [:lower:] [:upper:]) 
	echo -n " "$(lsblk --nodeps -no type /dev/sda)
	echo -n ", model" $(lsblk --nodeps -no model /dev/sda)
	echo ", s/n "$(lsblk --nodeps -no serial /dev/sda)
	echo -n "Press <Enter> if you wish to continue anyway (Control-C aborts) "
	read $ans
fi
if $(partprobe -ds /dev/sda | grep -q msdos) && [ ! -d /boot/efi ]; then
	boot="MBR"
elif [ ! -z "$(sgdisk -i1 /dev/sda | grep 'BIOS boot')" ] && [ ! -d /boot/efi ]; then
	boot="BIOS";
elif [ ! -z "$(sgdisk -i1 /dev/sda | grep 'EFI System')" ] && [ -d /boot/efi ]; then
	boot="EFI";
else
	echo "ERROR: Unrecognised boot scheme on \"Primary\" disk /dev/sda"
	exit 1
fi

# Now get the user to load the Secondary disk
if [ -b $disk ] && [ $ACCEPT_LOAD == 0 ] && ([ ! -O $lockf ] || [ "$(cat $lockf)" != "$(lsblk --nodeps -no serial $disk)" ]); then
	echo "ERROR:\"Secondary\" disk $disk must not be loaded before start of script"
	echo "(You may want to use the '-A' option to [A]llow if this is intended)"
	exit 1
fi
if [ ! -b $disk ]; then
	while [ ! -b $disk ]; do
		echo  -n -e "\rWaiting for disk $disk to be loaded, use 'Control-C' to safely abort ... "
		sleep 1
	done
	echo  -n -e "\r                                                                              "
	echo  -n -e "\rWaiting for disk $disk to settle ... "
	udevadm settle
	sleep 2
	echo $(lsblk --nodeps -no serial $disk) > $lockf
	echo -e "done.\n"
fi

# Perform sanity checks on secondary disk now that it's available
if [ "$(cat /sys/block/$devname/removable)" != "0" ]; then
	echo "WARNING: \"Secondary\" disk $disk is removable:"
	echo -n " $disk: "$(lsblk --nodeps -no tran $disk | tr [:lower:] [:upper:])
	echo -n " "$(lsblk --nodeps -no type $disk)
	echo -n ", model" $(lsblk --nodeps -no model $disk)
	echo ", s/n "$(lsblk --nodeps -no serial $disk)
	echo -n "Press <Enter> if you wish to continue anyway (Control-C aborts) "
	read $ans
fi
if mdadm --detail /dev/md0 | grep -q $disk; then
	echo "ERROR: \"Secondary\" disk $disk is still part of RAID /dev/md0!!"
	echo "(Perhaps you should actually be using 'recover_raid' instead)"
	exit 1
fi
if grep -q " $devname" /proc/mdstat; then
	echo "ERROR: \"Foreign\" RAID(s) detected on \"Secondary\" disk $disk"
	echo "(You should use blank_secondary to reset the second disk first)"
	exit 1
fi
if grep -q ^$disk /proc/mounts; then
	echo "ERROR: Some part of \"Secondary\" disk $disk is mounted!!"
	exit 1
fi
if [ $(cat /sys/block/md0/size) -gt $(cat /sys/block/$devname/size) ] ||
    ([ -b $part2 ] && [ $(cat /sys/block/md0/size) -gt $(cat /sys/block/$devname/"$devname"2/size) ]); then
	echo "ERROR: \"Secondary\" device $devname is too small to house RAID /dev/md0"
	exit 1
fi
if [ -b $part1 ] && [ ! -z "$(sgdisk -i1 $disk | grep 'GUID code')" ] &&
   (( [ $boot = "MBR" ] && ! $(partprobe -ds $disk | grep -q msdos) ) ||
    ( [ $boot = "BIOS" ] && [ -z "$(sgdisk -i1 $disk | grep 'BIOS boot')" ] ) ||
    ( [ $boot = "EFI" ] && [ -z "$(sgdisk -i1 $disk | grep 'EFI System')" ] )); then
	echo "ERROR: boot scheme on \"Secondary\" disk $disk does not match \"Primary\" disk /dev/sda"
	echo "(You should use blank_secondary to reset the second disk)"
	exit 1
fi

# Cross-compare state of the RAID partitions
if [ -b $part2 ] && [ "$(mdadm --examine -Y /dev/sda2 | grep MD_UUID | cut -d'=' -f2)" != "$(mdadm --examine -Y $part2 | grep MD_UUID | cut -d'=' -f2)" ]; then
	echo "ERROR: \"Secondary\" disk $part2 MD_UUID does not match \"Primary\" disk /dev/sda2"
	echo "(You should use blank_secondary to reset the second disk)"
	exit 1
fi
if [ -b $part2 ] && [ $(mdadm --examine -Y /dev/sda2 | grep TIME | cut -d'=' -f2) -lt $(mdadm --examine -Y $part2 | grep TIME | cut -d'=' -f2) ]; then
	echo "ERROR: \"Secondary\" disk $part2 TIME is ahead of \"Primary\" disk /dev/sda2"
	echo "(You should use blank_secondary to reset the second disk)"
	exit 1
fi
if [ -b $part2 ] && [ $(mdadm --examine -Y /dev/sda2 | grep EVENTS | cut -d'=' -f2) -lt $(mdadm --examine -Y $part2 | grep EVENTS | cut -d'=' -f2) ]; then
	echo "ERROR: \"Secondary\" disk $part2 EVENT count is higher than the \"Primary\" disk /dev/sda2"
	echo "(You should use blank_secondary to reset the second disk)"
	exit 1
fi

# Check if "Secondary" disk has been booted separately
if [ -b $part2 ] && [ "$(mdadm --examine $part2 | grep "Array State" | cut -d' ' -f7)" != "AA" ]; then
	echo "ERROR: \"Secondary\" disk $part2 has been used separately since last in sync"
	echo "(You should use blank_secondary to reset the second disk)"
	exit 1
fi

# Ensure the Primary disk is mounted on /boot/efi (ie. not Secondary) if applicable
if [ $boot = "EFI" ]; then
	umount /boot/efi
	mount /dev/sda1 /boot/efi
fi


# ------------- Everything looks sane so let's set to work ----------------

# First repartition the Secondary disk by cloning the Primary's table if blank
if [ "$(partprobe -s /dev/sda | cut -d' ' -f2-)" != "$(partprobe -s $disk | cut -d' ' -f2-)" ]; then
	if [ "$(partprobe -s $disk | cut -d' ' -f3-)" = "partitions" ]; then
		echo -n "Partitioning blank \"Secondary\" disk $disk ... "
		if [ $boot = "MBR" ]; then
			sfdisk -d /dev/sda | sfdisk /dev/sdb > /dev/null
		else
			sgdisk -R$disk /dev/sda > /dev/null
			sgdisk -G $disk > /dev/null
		fi
		partprobe $disk
		echo "done."
	else
		echo "ERROR: Partitioning of \"Secondary\" disk $disk does not match \"Primary\" disk /dev/sda"
		echo "(You should use blank_secondary to reset the second disk)"
		exit 1
	fi
fi

# Next install the bootloader to the Secondary disk
if [ $boot = "MBR" ] || [ $boot = "BIOS" ]; then
	echo -n "Installing GRUB bootloader in BIOS mode on \"Secondary\" disk $disk ... "
	grub-install $disk
	echo "done."
elif [ $boot = "EFI" ]; then
	# Ensure the primary disk is mounted on /boot/efi (not anything else)
	if ! grep "/boot/efi" /proc/mounts | grep -q /dev/sda1; then
		umount /boot/efi
		mount /dev/sda1 /boot/efi
	fi

	echo -n "Installing GRUB bootloader in UEFI mode on \"Secondary\" disk $disk ... "
	# Note Debian 9 `grub-install` essentially ignores the second argument and
	# copies the grub EFI shim to /boot/efi, regardless what is mounted there.
	# Instead, we refresh the grub install on the primary, and unmount it to 
	# ensure it is consistent, then copy the whole ESP to the second disk. 
	grub-install $disk 2> /dev/null		# suppress spurious warnings
	umount /boot/efi
	cp /dev/sda1 $part1
	# and remount the primary on /boot/efi again
	mount /dev/sda1 /boot/efi
	echo "done."
fi

# Then add the Secondary RAID partition as a replacement into the RAIO
echo "Adding \"Secondary\" disk $part2 to RAID array /dev/md0 ... "
echo "(This may safely be interrupted with 'Control-C' if you don't want to wait)"
mdadm --zero-superblock $part2 2> /dev/null	# will fail if disk is blank
mdadm /dev/md0 --add-spare $part2
sleep 1

# Check that the RAID is now rebuilding ie. the addition succeeded
if ! (mdadm --detail /dev/md0 | grep -q ", recovering $"); then
	echo "ERROR: RAID /dev/md0 failed to start rebuilding"
	exit 1
fi

# Now RAID is using both disks again, we delete the override interlock file
rm -f $lockf

# Finally report on recovery progress until it completes
while ! grep -q recovery /proc/mdstat; do 
	echo "Waiting for recovery ... "
	sleep 1;
done
#
while grep -q recovery /proc/mdstat; do 
	recovery=$(grep recovery /proc/mdstat | cut -c32-)
	echo -e -n "\r$recovery"
	sleep 1;
done
echo -e "\ndone."
